#include "asmgen.hpp"
#include <cassert>

Registers rv32regs;

std::unique_ptr<std::string>
to_riscv32_asm(std::unique_ptr<std::string>& koopa_ir_text)
{
    std::ostringstream decl;
    std::ostringstream inst;

    decl << "\t.text\n";

    visit_koopa_bin(koopa_ir_text->c_str(), decl, inst);

    return std::make_unique<std::string>(decl.str() + inst.str());
}

void visit_koopa_bin(const char*         str,
                     std::ostringstream& decl,
                     std::ostringstream& inst)
{
    koopa_program_t    program;
    koopa_error_code_t ret = koopa_parse_from_string(str, &program);
    assert(ret == KOOPA_EC_SUCCESS);

    koopa_raw_program_builder_t builder = koopa_new_raw_program_builder();
    koopa_raw_program_t         raw = koopa_build_raw_program(builder, program);
    koopa_delete_program(program);

    // something todo with raw

    visit_koopa(raw, decl, inst);

    koopa_delete_raw_program_builder(builder);
}

void visit_koopa(const koopa_raw_program_t& program,
                 std::ostringstream&        decl,
                 std::ostringstream&        inst)
{
    for (unsigned i = 0; i < program.values.len; i++) {
        visit_koopa(
            reinterpret_cast<koopa_raw_value_t>(program.values.buffer[i]), decl,
            inst);
    }

    for (unsigned i = 0; i < program.funcs.len; i++) {
        visit_koopa(
            reinterpret_cast<koopa_raw_function_t>(program.funcs.buffer[i]),
            decl, inst);
    }
}

void visit_koopa(const koopa_raw_value_t& value,
                 std::ostringstream&      decl,
                 std::ostringstream&      inst)
{
    const auto& kind = value->kind;

    switch (kind.tag) {
        case KOOPA_RVT_RETURN: visit_koopa(kind.data.ret, decl, inst); break;

        case KOOPA_RVT_BINARY:
            visit_koopa(kind.data.binary, decl, inst, &kind);
            break;

        default: assert(false); break;
    }
}

void visit_koopa(const koopa_raw_return_t& ret,
                 std::ostringstream&       decl,
                 std::ostringstream&       inst)
{
    // std::cout << "Encountered return expression.\n";

    // 检查有无返回值, 如果有则默认i32, 使用li加载到a0后ret, 如无直接ret

    if (ret.value != nullptr) {
        auto value = ret.value;
        switch (value->kind.tag) {
            case KOOPA_RVT_INTEGER:
                inst << "\tli a0, " << value->kind.data.integer.value << '\n';
                break;

            case KOOPA_RVT_BINARY: {
                int reg =
                    rv32regs.store(&value->kind, Registers::StorageType::ret);
                int regfrom = rv32regs.find(&value->kind);
                inst << "\tmv\t" << rv32regs.name(reg) << ", "
                     << rv32regs.name(regfrom) << '\n';
            } break;
            default:
                // no other choice now
                assert(false);
        }
    }

    inst << "\t"
         << "ret" << '\n';
}

void visit_koopa(const koopa_raw_function_t& func,
                 std::ostringstream&         decl,
                 std::ostringstream&         inst)
{
    // std::cout << "In func " << func->name << ": \n";
    // skip first char '@'
    decl << "\t.globl " << func->name + 1 << '\n';

    inst << func->name + 1 << ":\n";
    for (unsigned i = 0; i < func->bbs.len; i++) {
        visit_koopa(
            reinterpret_cast<koopa_raw_basic_block_t>(func->bbs.buffer[i]),
            decl, inst);
    }
}

void visit_koopa(const koopa_raw_basic_block_t& block,
                 std::ostringstream&            decl,
                 std::ostringstream&            inst)
{
    // std::cout << "In basic block " << block->name << ": \n";

    for (unsigned i = 0; i < block->params.len; i++) {
        visit_koopa(*(reinterpret_cast<const koopa_raw_slice_t*>(
                        block->params.buffer[i])),
                    decl, inst);
    }

    for (unsigned i = 0; i < block->used_by.len; i++) {
        visit_koopa(*(reinterpret_cast<const koopa_raw_slice_t*>(
                        block->used_by.buffer[i])),
                    decl, inst);
    }

    for (unsigned i = 0; i < block->insts.len; i++) {
        // ?
        visit_koopa(reinterpret_cast<koopa_raw_value_t>(block->insts.buffer[i]),
                    decl, inst);
    }
}

void visit_koopa(const koopa_raw_slice_t& slice,
                 std::ostringstream&      decl,
                 std::ostringstream&      inst)
{
    for (unsigned i = 0; i < slice.len; i++) {
        auto ptr = slice.buffer[i];
        switch (slice.kind) {
            case KOOPA_RSIK_FUNCTION:
                visit_koopa(reinterpret_cast<koopa_raw_function_t>(ptr), decl,
                            inst);
                break;
            case KOOPA_RSIK_BASIC_BLOCK:
                visit_koopa(reinterpret_cast<koopa_raw_basic_block_t>(ptr),
                            decl, inst);
                break;
            case KOOPA_RSIK_VALUE:
                visit_koopa(reinterpret_cast<koopa_raw_value_t>(ptr), decl,
                            inst);
                break;

            default: assert(false); break;
        }
    }
}

void visit_koopa(const koopa_raw_binary_t&     bin,
                 std::ostringstream&           decl,
                 std::ostringstream&           inst,
                 const koopa_raw_value_kind_t* kindptr)
{
    auto get_reg_num = [&](const koopa_raw_value_kind_t* _kdp) {
        int reg;
        if ((reg = rv32regs.find(_kdp)) != Registers::None) {
            return reg;
        }
        else {
            reg = rv32regs.store(_kdp, Registers::StorageType::tmp);
            switch (_kdp->tag) {
                // 如果是立即数, 就先加载到寄存器
                case KOOPA_RVT_INTEGER: {
                    inst << "\tli\t" << rv32regs.name(reg) << ", "
                         << _kdp->data.integer.value << '\n';
                } break;
                default: break;
            }
            return reg;
        }
    };
    int lhsreg = get_reg_num(&bin.lhs->kind),
        rhsreg = get_reg_num(&bin.rhs->kind);

    switch (bin.op) {
        case KOOPA_RBO_EQ: {
            // int lhsreg = rv32regs.store(kindptr,
            // Registers::StorageType::tmp);
            // inst << "\tli\t" << rv32regs.name(lhsreg) << "\t, "
            //      << bin.lhs->kind.data.integer.value << '\n';

            // int rhsreg = rv32regs.store(kindptr,
            // Registers::StorageType::tmp);
            // inst << "\tli\t" << rv32regs.name(rhsreg) << "\t, "
            //      << bin.rhs->kind.data.integer.value << '\n';

            inst << "\txor\t" << rv32regs.name(lhsreg) << ", "
                 << rv32regs.name(lhsreg) << ", " << rv32regs.name(rhsreg)
                 << "\n";
            // 根据实例, 同一条指令的源寄存器和目的寄存器可以相同
            // 考虑到不存在多线程/多进程, 可以直接在此释放两个
            // 并获取一个新的存放结果, 实际就是释放一个
            // rv32regs.release(rhsreg);
            inst << "\tseqz\t" << rv32regs.name(lhsreg) << ", "
                 << rv32regs.name(lhsreg) << "\n";
            rv32regs.store(kindptr, lhsreg);

        } break;
        case KOOPA_RBO_SUB: {
            inst << "\tsub\t" << rv32regs.name(lhsreg) << ", "
                 << rv32regs.name(lhsreg) << ", " << rv32regs.name(rhsreg)
                 << "\n";
            rv32regs.store(kindptr, lhsreg);
        } break;
        default: assert(false); break;
    }
}